
// Copyright (c) 2007, Nils Jonas Norberg, See COPYING.txt for details

#include <stdio.h>
#include <stdlib.h>

#include "psdlite.h"

namespace psdlite
{
	struct buffered_file
	{
		buffered_file( const char* fname )
			:iter_(0)
		{
			mem_.clear();
			FILE* f = fopen( fname, "rb" );
			if (!f)
				return;

			fseek( f,0,SEEK_END );
			mem_.resize( ftell(f) );
			fseek( f, 0, SEEK_SET );
			fread( &mem_[0], mem_.size(), 1, f );
			fclose( f );
		}

		size_t get_pos(){ return iter_; }
		void set_pos( size_t pos)
		{
			iter_ = pos;
			if ( iter_ > mem_.size() )
			{
				throw error_code_invalid_file;
			}
		}

		void pad_even()
		{
			iter_ = (iter_ + 1) & ~1;
			if (iter_ > mem_.size())
			{
				throw error_code_invalid_file;
			}
		}

		u32 getu32()
		{
			if (iter_ + 4 > mem_.size() )
			{
				throw error_code_invalid_file;
			}

			u8* p = (u8*)&mem_[iter_];
			iter_ += 4;
			return (p[0] << 24) | (p[1] << 16) | (p[2] << 8) | p[3];
		}

		u16 getu16()
		{
			if (iter_ + 2 > mem_.size() )
			{
				throw error_code_invalid_file;
			}

			u8* p = (u8*)&mem_[iter_];
			iter_ += 2;
			return (p[0] << 8) | p[1];
		}

		s16 gets16()
		{
			return (s16)getu16();
		}

		u8 getu8()
		{
			if (iter_ + 1 > mem_.size() )
			{
				throw error_code_invalid_file;
			}

			u8* p = (u8*)&mem_[iter_];
			iter_ += 1;
			return p[0];
		}

		s8 gets8()
		{
			return (s8)getu8();
		}

		void skip( u32 bytes )
		{
			iter_ += bytes;

			if (iter_ > mem_.size() )
			{
				throw error_code_invalid_file;
			}
		}

		void skip_pstring()
		{
			u8 s = getu8();
			skip( s );
			pad_even();
		}

		void get_pstring(std::string& str )
		{
			u8 s = getu8();
			iter_ += s;
			if (iter_ > mem_.size() )
			{
				throw error_code_invalid_file;
			}

			str.assign( &mem_[ iter_ - s ], &mem_[iter_] );

			pad_even();
		}

	protected:
		std::vector<s8> mem_;
		size_t iter_;
	};

	/////////////////////////////////////////////////////////////

	struct loader
	{
		loader(buffered_file& file)
			:file_(file)
		{
		}

	private:
		buffered_file& file_;

		void operator=(const loader&){}; // no assignement operator

		void parse_header( layered_image& dest )
		{
			u32 sig = file_.getu32();
			(void)sig;

			u16 ver = file_.getu16();
			(void)ver;

			file_.skip( 6 );

			u16 channels = file_.getu16();
			(void)channels;

			u32 rows = file_.getu32();
			u32 columns = file_.getu32();

			dest.size_.set( columns, rows );

			u16 depth = file_.getu16();
			if (depth != 8)
			{
				throw error_code_not_supported; // only support 8bit depth
			}

			u16 mode = file_.getu16();

			if ( mode != 3 )
			{
				throw error_code_not_supported; // only support RGB mode
			}
		}

		void skip_block()
		{
			u32 size = file_.getu32();
			file_.skip( size );
		}

		void parse_image_resource_block( layered_image& /*dest*/)
		{
			u32 block_type = file_.getu32();
			(void)block_type;

			u16 block_id = file_.getu16();
			(void)block_id;

			file_.skip_pstring();

			u32 size = file_.getu32();
			file_.skip( size );
		}

		void parseRAWChannel( bitmap& dest, int color_channel)
		{
			vi2 s = dest.get_size();

			for ( int y = 0 ; y != s.y ; ++y)
			{
				for ( int x = 0 ; x != s.x ; ++x )
				{
					u8 v = file_.getu8();
					dest.set_single_channel(x, y, color_channel, v );
				}
			}
		}

		void parseRLEChannel( bitmap& dest, int color_channel )
		{
			// RLE compression...
			vi2 s = dest.get_size();

			std::vector<u16> scanline_byte_counts( s.y );

			// get bytecounts for all scanlines
			for ( int y = 0 ; y != s.y ; ++y )
			{
				scanline_byte_counts[y] = file_.getu16();
			}

			for ( int y = 0 ; y != s.y ; ++y )
			{
				int line_bytes = scanline_byte_counts[y];
				for ( int x = 0 ; line_bytes && (x < s.x) ; )
				{
					int control_byte = file_.gets8();
					--line_bytes;
					if (control_byte < 0)
					{
						int count = 1 - control_byte;

						// RLE
						if (line_bytes)
						{
							s8 v = file_.getu8();
							--line_bytes;

							for ( ; count ; --count, ++x )
								dest.set_single_channel(x, y, color_channel, v );
						}
					}
					else
					{
						int count = 1 + control_byte;

						// RAW
						for ( ; count && line_bytes; --line_bytes, --count, ++x )
							dest.set_single_channel(x, y, color_channel, file_.getu8() );
					}
				}
			}
		}

		void parse_layer_channel_data( bitmap& dest)
		{
			int cc = dest.get_channel_count();
			for ( int channel = 0 ; channel != cc ; ++channel )
			{
				int color_channel = channel;

				// if only RGB skip A in ARGB
				if (cc == 3)
				{
					color_channel++;
				}

				u16 compression = file_.getu16();

				switch (compression)
				{
				case 0: // raw data
					parseRAWChannel( dest, color_channel );
					break;
				case 1: // rle.. good
					parseRLEChannel( dest, color_channel );
					break;
				case 2:
				case 3:
					throw error_code_not_supported;
				}
			}
		}

		void parse_layer_pixel_data( layered_image& dest )
		{
			for ( u32 i = 0 ; i != dest.layers_.size() ; ++i )
				parse_layer_channel_data( dest.layers_[i].data_ );
		}

		void parse_layer_record( layered_image& dest )
		{
			u32 top = file_.getu32();
			u32 left = file_.getu32();
			u32 bottom = file_.getu32();
			u32 right = file_.getu32();

			u16 channel_count = file_.getu16();
			file_.skip( 6 * channel_count ); // skip channel length info

			u32 blend_mode_sig = file_.getu32();
			(void)blend_mode_sig;

			u32 blend_mode_key = file_.getu32();
			(void)blend_mode_key;

			u8 opacity = file_.getu8();
			(void)opacity;

			u8 clipping = file_.getu8();
			(void)clipping;

			u8 flags = file_.getu8();
			(void)flags;

			u8 filler = file_.getu8();
			(void)filler;

			u32 extra_size = file_.getu32();
			size_t endpos = file_.get_pos() + extra_size;

			skip_block(); // layer mask adjustment data
			skip_block(); // layer blending ranges

			std::string l_name;
			file_.get_pstring( l_name );

			// jump past adjustment layer info
			file_.set_pos( endpos );

			// add layer to dest
			dest.layers_.push_back( layer() );
			layer& l = dest.layers_.back();
			l.name_ = l_name;
			l.offs_.set( left, top );
			l.data_.resize( right - left, bottom - top );
			l.data_.set_channel_count( channel_count );
			l.opacity_ = opacity;
			l.flags_ = flags;
		}

		void parse_layer_structure( layered_image& dest )
		{
			s16 l_count = (s16)abs( file_.gets16() );

			for (int i = 0 ; i != l_count ; ++i )
				parse_layer_record( dest );

			file_.pad_even();
		}

		void parse_layer_info( layered_image& dest )
		{
			u32 size = file_.getu32();

			size_t endpos = file_.get_pos() + size;

			parse_layer_structure( dest );
			parse_layer_pixel_data( dest );

			file_.set_pos( endpos );
		}

		void parse_layer_and_mask( layered_image& dest )
		{
			u32 size = file_.getu32();
			size_t endpos = file_.get_pos() + size;

			parse_layer_info( dest );
			skip_block(); // parse_mask_info( dest, ptr, stop );

			file_.set_pos( endpos );
		}

	public:
		void parse_layered_image( layered_image& dest )
		{
			parse_header( dest );
			skip_block(); // parse_color_data( dest );
			skip_block(); // parse_image_resources( dest );
			parse_layer_and_mask( dest );

			// skip composite image...
		}
	};

	error_code load_layered_image( layered_image& dest, const char* fname )
	{
		try
		{
			// clear dest
			dest.layers_.clear();
			dest.size_.set(0,0);

			// load file
			buffered_file file( fname );

			loader l( file );
			l.parse_layered_image( dest );
		}
		catch ( error_code e )
		{
			return e;
		}
		catch ( ... )
		{
			return error_code_invalid_file;
		}

		return error_code_no_error;
	}
}
